# question2table.py
# usage: 
# python ./src/vector/question2table.py

import os,sys
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, project_root)
from config import project_config
import pandas as pd
import json
from query.extract_tables import extract_tables
from tools.clean_question import clean_question

def main():
    # 读取 samples.jsonl
    samples_df = pd.read_json(project_config.samples_path, lines=True, encoding='utf-8')

    # 表 -> 问题列表（允许重复，后续去重）
    table_to_questions = {}

    for _, row in samples_df.iterrows():
        question = row["question"]
        sql = row["sql"]

        cleaned_q = clean_question(question)
        tables = extract_tables(sql)

        if not tables:
            print(f"警告：SQL 未提取到表名：{sql}")
            continue

        for table in tables:
            if table not in table_to_questions:
                table_to_questions[table] = []
            table_to_questions[table].append(cleaned_q)

    # ✅ 关键：在输出前对每个表的问题列表进行【去重】
    output_data = []
    for table, questions in table_to_questions.items():
        unique_questions = list(dict.fromkeys(questions))  # 保持顺序去重
        item = {
            "table": [table],
            "questions": unique_questions  # ✅ 确保是去重后的列表
        }
        output_data.append(item)

    # 写入文件，每行一个 JSON 对象
    with open("tmp_question2sql.json", "w", encoding="utf-8") as f:
        for item in output_data:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")

    print("✅ 已成功生成 tmp_question2sql.json，问题已去重")


if __name__ == "__main__":
    main()